import os
from typing import Dict, List, Optional

from langchain_community.llms import ChatGLM
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_openai import ChatOpenAI


# 1. 初始化本地 ChatGLM3-6B 模型
def init_chatglm_model():
    return ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-plus",  # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    # other params...
)
# 2. 初始化 Tavily 搜索工具
def init_tavily_tool():
    return TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))


# 3. 创建聊天代理
def create_chat_agent(llm, tools):
    prompt = ChatPromptTemplate.from_messages([
        SystemMessage(content="你是一个有帮助的AI助手，可以访问网络搜索信息"),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad")
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    return AgentExecutor(agent=agent, tools=tools, verbose=True)


# 4. 主聊天循环
def chat_loop(agent_executor):
    chat_history = []

    print("欢迎使用 ChatGLM3-6B 助手 (输入 'exit' 退出)")
    while True:
        try:
            user_input = input("\n你: ")
            if user_input.lower() == 'exit':
                break

            # 执行对话
            result = agent_executor.invoke({
                "input": user_input,
                "chat_history": chat_history
            })

            # 显示结果并更新历史
            print(f"\n助手: {result['output']}")
            chat_history.extend([
                HumanMessage(content=user_input),
                AIMessage(content=result['output'])
            ])

        except KeyboardInterrupt:
            print("\n退出聊天...")
            break
        except Exception as e:
            print(f"发生错误: {str(e)}")


# 5. 主函数
def main():
    # 初始化组件
    llm = init_chatglm_model()
    tavily_tool = init_tavily_tool()

    # 创建代理执行器
    agent_executor = create_chat_agent(llm, [tavily_tool])

    # 启动聊天循环
    chat_loop(agent_executor)


if __name__ == "__main__":
    # 环境变量检查
    if not os.getenv("TAVILY_API_KEY"):
        print("请先设置 TAVILY_API_KEY 环境变量")
        exit(1)

    main()